#include "UCTSearch.hpp"

/***********************************************/
/********************UCTNode********************/
/***********************************************/
UCTNode::UCTNode() {}

UCTNode::UCTNode(UCTNode *parent_, int **board_, int *top_, int M_, int N_, int last_x_, int last_y_, Role role_)
    : parent(parent_), visited_num(0), M(M_), N(N_), gain(0), last_x(last_x_), last_y(last_y_),
      terminal_flag(NOT_DEFINE), role(role_)
{
    /*copy board and top*/
    board = new int *[M_];
    for (int i = 0; i < M_; ++i)
    {
        board[i] = new int[N_];
        std::copy(board_[i], board_[i] + N_, board[i]);
    }
    top = new int[N];
    std::copy(top_, top_ + N_, top);
}

/****************************************
 * allocate()
 * Function:
 *      分配内存，在初始化使用
 *      这里直接分配最大的棋盘大小，放置反复
 *      进行内存的释放
 ****************************************/
void UCTNode::allocate()
{
    board = new int *[ROW_MAX_NUMBER];
    for (int i = 0; i < ROW_MAX_NUMBER; ++i)
    {
        board[i] = new int[COL_MAX_NUMBER];
    }
    top = new int[COL_MAX_NUMBER];
}

/****************************************
 * reset
 * Function:
 *      用来设置或者重置节点的信息，
 *      不会重新分配内存
 ****************************************/
void UCTNode::reset(UCTNode *parent_, int **board_, int *top_, int M_, int N_, int last_x_, int last_y_, Role role_)
{
    M = M_;
    N = N_;
    parent = parent_;
    child.clear();
    possible_pos.clear();
    visited_num = 0;
    gain = 0;
    last_x = last_x_;
    last_y = last_y_;
    terminal_flag = NOT_DEFINE;
    role = role_;
    /*copy board info*/
    for (int i = 0; i < M; ++i)
        std::copy(board_[i], board_[i] + N, board[i]);
    std::copy(top_, top_ + N, top);
}

UCTNode::~UCTNode()
{
    delete[] top;
    for (int i = 0; i < ROW_MAX_NUMBER; ++i)
        delete[] board[i];
    delete[] board;
}

/****************************************
 * set_possible_pos
 * Function:
 *      设置节点可下的列，用在走子之后调用
 ****************************************/
void UCTNode::set_possible_pos()
{
    /*初始化可下的列*/
    for (int i = 0; i < N; ++i)
    {
        if (top[i] > 0)
            possible_pos.push_back(i);
    }
}

void UCTNode::copy_board(int **copy_board)
{
    for (int i = 0; i < M; ++i)
    {
        std::copy(board[i], board[i] + N, copy_board[i]);
    }
}

void UCTNode::copy_top(int *copy_top) { std::copy(top, top + N, copy_top); }

/**************************************************/
/********************UCT search********************/
/**************************************************/
UCTSearch::UCTSearch(int M_, int N_, int noX, int noY)
    : M(M_), N(N_), no_x(noX), no_y(noY), generator(SEED), dis(0, INT_MAX), pool_p(0)
{
    /*allocate memory cache*/
    board_cache = new int *[ROW_MAX_NUMBER];
    for (int i = 0; i < ROW_MAX_NUMBER; ++i)
        board_cache[i] = new int[COL_MAX_NUMBER];
    top_cache = new int[COL_MAX_NUMBER];
    /*allocate memory pool*/
    node_pool = new UCTNode[NODE_NUMBER];
    for (int i = 0; i < NODE_NUMBER; ++i)
        node_pool[i].allocate();
}

UCTSearch::~UCTSearch()
{
    // delete cache
    delete[] top_cache;
    for (int i = 0; i < ROW_MAX_NUMBER; ++i)
        delete[] board_cache[i];
    delete[] board_cache;
    // delete memory pool
    delete[] node_pool;
}

/****************************************
 * reset
 * Function:
 *      开始或者重新设置搜索树，只需要简单更改根节点即可
 ****************************************/
void UCTSearch::reset(int M_, int N_, int **board_, int *top_, int no_x_, int no_y_)
{
    M = M_;
    N = N_;
    no_x = no_x_;
    no_y = no_y_;
    //重新设置node
    // assert(pool_p == 0);
    pool_p = 0;
    root = &node_pool[pool_p++];
    root->reset(nullptr, board_, top_, M, N, -1, -1, COMPUTER);
    generator.seed(SEED);
}

/**************************************************
 *上面是一系列准备工作
 *下面是和UCTSearch相关的搜索函数
 **************************************************/
UCTNode *UCTSearch::expand(UCTNode *v)
{
    /*随机选择一步*/
    int random_rank = dis(generator) % v->possible_pos.size();
    int next_col = v->possible_pos[random_rank]; //要走的一列
    int next_row = v->top[next_col] - 1;         //要走的一行
    /*在可行的走法中去掉当前这一步*/
    std::swap(v->possible_pos[random_rank], v->possible_pos.back());
    v->possible_pos.pop_back();
    /*扩展当前的节点v*/
    Role next_role = (v->role == USER) ? COMPUTER : USER;
    int chess_number = (v->role == USER) ? 1 : 2;
    UCTNode *child_p = &node_pool[pool_p++];
    if (pool_p >= NODE_NUMBER)
        std::cerr << " pool_p = " << pool_p << endl;
    child_p->reset(v, v->board, v->top, M, N, next_row, next_col, next_role);
    v->child.push_back(child_p);
    /*在后代中体现走的这一步*/
    put_chess(next_row, next_col, child_p->board, child_p->top, chess_number);
    /*Init child*/
    child_p->set_possible_pos();
    return child_p;
}

bool UCTSearch::is_terminal(UCTNode *v)
{
    if (v->terminal_flag != NOT_DEFINE) //只要不是NOT_DEFINE，那么直接返回
        return v->terminal_flag;
    if (isTie(N, v->top))
    {
        v->terminal_flag = TIE_TERMINAL;
        return true;
    }
    /*注意，下面的v的role和判断函数是反过来的，因为上一步是相反的角色走的棋*/
    if (v->role == USER && machineWin(v->last_x, v->last_y, M, N, v->board))
    {
        v->terminal_flag = WIN_TERMINAL;
        return true;
    }
    if (v->role == COMPUTER && userWin(v->last_x, v->last_y, M, N, v->board))
    {
        v->terminal_flag = WIN_TERMINAL;
        return true;
    }
    v->terminal_flag = NOT_TERMINAL; //表明不是终点
    return false;
}

UCTNode *UCTSearch::best_child(UCTNode *p, double coef)
{
    double max_value = -DBL_MAX;
    int max_rank;
    for (int i = 0; i < p->child.size(); ++i)
    {
        const UCTNode &child = *p->child[i];
        double val = child.gain / child.visited_num + coef * sqrt(2.0 * log(p->visited_num) / child.visited_num);
        if (val > max_value)
        {
            max_value = val;
            max_rank = i;
        }
    }
    return p->child[max_rank];
}

/*tree_policy一直循环，直到找到第一个可扩展的节点并扩展后返回，或者到达终点*/
UCTNode *UCTSearch::tree_policy(UCTNode *node_ptr)
{
    /*这里不得不使用指针，因为需要不断更新node_ptr*/
    /*先判断是否到达终点*/
    while (!is_terminal(node_ptr))
    {
        if (!node_ptr->possible_pos.empty())
        {
            return expand(node_ptr);
        }
        else
        {
            node_ptr = best_child(node_ptr, COEF);
        }
    }
    return node_ptr;
}

double UCTSearch::default_policy(UCTNode *v)
{
    if (v->terminal_flag > NOT_TERMINAL)
    {
        // v->terminal_flag只能是WIN或者TIE，那么不用再模拟了
        return (v->terminal_flag == WIN_TERMINAL) ? WIN_GAIN : TIE_GAIN;
    }
    /*执行模拟，先拷贝当前状态*/
    assert(v->visited_num == 0);
    v->copy_board(board_cache);
    v->copy_top(top_cache);
    Role role = v->role; //首先是v代表的一方下棋
    int chess_number = (role == USER) ? 1 : 2;
    double final_gain = 0;

    vector<int> candidate_pos(v->possible_pos.begin(), v->possible_pos.end());
    while (!candidate_pos.empty())
    {
        /*随机选择一列落子*/
        int random_rank = dis(generator) % candidate_pos.size();
        int rand_col = candidate_pos[random_rank];
        int row = top_cache[rand_col] - 1;
        /*更新棋盘*/
        put_chess(row, rand_col, board_cache, top_cache, chess_number);
        /*更新候选位置*/
        if (top_cache[rand_col] <= 0)
        {
            std::swap(candidate_pos[random_rank], candidate_pos.back());
            candidate_pos.pop_back();
        }
        /*判断胜负并更新走棋方*/
        if (isTie(N, top_cache))
        {
            final_gain = TIE_GAIN;
            break;
        }
        if (role == USER)
        {
            if (userWin(row, rand_col, M, N, board_cache))
            {
                /*如果v(开始模拟的)这一步是USER，表明上一步是COMPUTER走的，因此判负*/
                final_gain = (v->role == USER) ? LOSE_GAIN : WIN_GAIN;
                break;
            }
            else
            {
                role = COMPUTER;
                chess_number = 2;
            }
        }
        else // role==COMPUTER
        {
            if (machineWin(row, rand_col, M, N, board_cache))
            {
                /*如果v这一步是USER，表明上一步是COMPUTER走的，因此判胜*/
                final_gain = (v->role == USER) ? WIN_GAIN : LOSE_GAIN;
                break;
            }
            else
            {
                role = USER;
                chess_number = 1;
            }
        }
    }
    return final_gain;
}

void UCTSearch::back_prop(UCTNode *node_ptr, double delta_gain)
{
    //当node_ptr不为空，则一直向上传递
    while (node_ptr)
    {
        node_ptr->visited_num++;
        node_ptr->gain += delta_gain;
        node_ptr = node_ptr->parent;
        delta_gain = - delta_gain;
    }
}

void UCTSearch::put_chess(int x, int y, int **board, int *top, int chess_number)
{
    board[x][y] = chess_number;
    if (no_y == y && no_x == x - 1)
        top[y] = x - 1;
    else
        top[y] = x;
}

int UCTSearch::search(clock_t start)
{
    root->set_possible_pos();
    root->terminal_flag = NOT_TERMINAL;
    UCTNode *v = nullptr;
    double delta_gain;
    int tie_number = 0;
    while (((double)clock() - start) / CLOCKS_PER_SEC < TIME_LIMIT)
    {
        /*首先找到一个可以扩展的节点，并将其扩展*/
        v = tree_policy(root);
        /*对扩展节点做一次模拟*/
        delta_gain = default_policy(v);
        if (delta_gain == TIE_GAIN)
            tie_number += 1;
        /*模拟的增益回传*/
        back_prop(v, delta_gain);
    }
    if (root->visited_num == 0)
    {
        clean();
        return 5;
    }
    auto child = best_child(root, 0);
    int best_col = child->last_y;
    std::cerr << "Simulate number : " << root->visited_num << ". Node number = " << pool_p
              << ". Tie number = " << tie_number << endl;
    clean();
    return best_col;
}

/****************************************
 * clean
 * Function:
 *      决策后用来清理数据的，减少内存开销
 ****************************************/
void UCTSearch::clean() { pool_p = 0; }